from typing import TypedDict, Annotated
import requests
from dotenv import load_dotenv
import os
import logging
from util import writeFile

load_dotenv()

class NovelState(TypedDict):
    current_outline: str
    chapters: dict
    refinement_count: int
    current_chapter: str

class NovelNodes:
    def __init__(self):
        self.endpoint = 'https://open.bigmodel.cn/api/paas/v4/chat/completions'
        self.api_key = '734c6ee4a5524052a1fc44580660dc38.F5A02rjHWxocAsaT'
        self.model = 'glm-4-flash'

    def _call_llm(self, prompt: str, max_tokens=500):
        headers = {
            'Authorization': f'Bearer {self.api_key}',
            'Content-Type': 'application/json'
        }
        data = {
            "model": self.model,
            "messages": [{"role": "user", "content": prompt}],
            "max_tokens": max_tokens
        }
        
        logging.info(f"调用LLM API：提示词长度{len(prompt)}字符")
        response = requests.post(self.endpoint, json=data, headers=headers)
        return response.json()['choices'][0]['message']['content']

    def generate_initial_outline(self, state: NovelState):
        prompt = "作为小说创作助手，请生成包含5个章节的小说大纲，包含主要人物和情节发展。"
        outline = self._call_llm(prompt)
        writeFile('output/outline.txt', outline)
        return {"current_outline": outline, "refinement_count": 0}

    def refine_outline(self, state: NovelState):
        prompt = f"当前大纲：{state['current_outline']}\n请根据以下反馈改进大纲：{state.get('feedback', '')}"
        new_outline = self._call_llm(prompt)
        return {"current_outline": new_outline, "refinement_count": state['refinement_count'] + 1}

    def write_chapter(self, state: NovelState, chapter_name: str):
        prompt = f"小说大纲：{state['current_outline']}\n请撰写章节《{chapter_name}》，包含场景描写和人物对话，控制在500字左右。"
        content = self._call_llm(prompt)[:50] + "..."  # 截断前50字符
        
        filename = f"output/chapters/{chapter_name}.txt"
        writeFile(filename, content)
        return {
            "chapters": {**state['chapters'], chapter_name: content},
            "word_count": len(content)
        }

    def final_review(self, state: NovelState):
        prompt = f"完整审查小说内容：{state['chapters']}\n请给出最终修改建议。"
        return {"review": self._call_llm(prompt)}